# test_rsa_oaep.py

from rsa_core import RSA, DEFAULT_HASH_FUNC



if __name__ == '__main__':
    # --- 用 RSA 类替代原来的函数调用 ---
    key_bits = 4096  # 标准位数
    print(f"开始生成 {key_bits} 位 RSA 密钥对...\n")
    rsa = RSA(bits=key_bits, hash_func=DEFAULT_HASH_FUNC)

    public_key = rsa.public_key
    private_key = rsa.private_key

    print("\n--- 密钥生成完毕 ---")
    print(f"公钥 (n, e): n={public_key[0]}, e={public_key[1]}")
    # 私钥 d 一般不直接打印

    # 要加密的消息
    message_str_short = "Hello OAEP!"
    message_bytes_short = message_str_short.encode('utf-8')

    # 生成一个稍长的测试消息
    message_str_long = (
        "这是一个用于测试RSA-OAEP的相对较长的消息。希望它能正常工作！" * 2
    )
    message_bytes_long = message_str_long.encode('utf-8')
    print(f"长消息字节长度: {len(message_bytes_long)}")

    # 可选标签
    label = b"MyApplication-Context-Specific-Label"

    messages_to_test = {
        "短消息": message_bytes_short,
        "长消息": message_bytes_long,
    }

    for name, msg in messages_to_test.items():
        print(f"\n--- 测试: {name} ---")
        print(
            f"原始消息: {msg[:60]}..."
            if len(msg) > 60
            else msg
        )
        # 加密
        print("--- 开始加密 (OAEP) ---")
        try:
            ciphertext = rsa.encrypt(msg, label=label)
            print(f"密文 (整数): {ciphertext}\n")

            # 解密
            print("--- 开始解密 (OAEP) ---")
            decrypted = rsa.decrypt(ciphertext, label=label)

            if decrypted == msg:
                print(f"成功！{name} 加解密一致。")
            else:
                print(f"失败！{name} 解密结果不一致。")
                print(f"解密后: {decrypted}")

        except ValueError as e:
            # OAEP 长度或其它错误
            n = public_key[0]
            k_bytes = (n.bit_length() + 7) // 8
            h_len = DEFAULT_HASH_FUNC().digest_size
            max_len = k_bytes - 2 * h_len - 2
            print(f"错误: {e}")
            print(f"  (k_bytes={k_bytes}, hash_len={h_len}, max_msg_len={max_len}, got={len(msg)})")

    # 测试消息过长
    print("\n--- 测试消息过长 ---")
    n = public_key[0]
    k_bytes = (n.bit_length() + 7) // 8
    h_len = DEFAULT_HASH_FUNC().digest_size
    max_len = k_bytes - 2 * h_len - 2
    too_long = b"A" * (max_len + 1)

    try:
        rsa.encrypt(too_long, label=label)
    except ValueError as e:
        print(f"预期错误 (消息过长): {e}")

    # 测试解密错误：错误标签
    print("\n--- 测试解密错误 (错误标签) ---")
    ct = rsa.encrypt(message_bytes_short, label=label)
    wrong_label = b"Wrong-Label"
    try:
        rsa.decrypt(ct, label=wrong_label)
    except ValueError as e:
        print(f"预期错误 (错误标签): {e}")

    # 测试解密错误：篡改后的密文
    print("\n--- 测试解密错误 (篡改密文) ---")
    tampered = ct + 1
    try:
        rsa.decrypt(tampered, label=label)
        print("意外：篡改密文解密成功（概率极低）")
    except ValueError as e:
        print(f"预期错误 (篡改密文): {e}")
